import sys
import os
import matplotlib.pyplot as plt
import codecs

def LogLoss(log_file):
    losses = []
    with codecs.open(log_file, 'r', encoding=u'utf-8', errors='ignore') as fr:
        for line in fr:
            if "'Loss':" in line:
                loss = float(line.split(":")[1].strip("}\n"))
                losses.append(loss)
    plt.ylabel("loss")
    plt.xlabel("step")
    plt.title("training loss on NLI")
    plt.plot(range(len(losses)), losses)
    plt.savefig(f"./SentCluster_Loss.png")

def FewShotPerf(log_dir):
    log_files = [fname for fname in os.listdir(log_dir) if ".log" in fname]
    for filename in log_files:
        log_file = os.path.join(log_dir, filename)
        fewshot_losses = []
        fewshot_accs = []
        with open(log_file, 'r') as fr:
            for line in fr:
                if "####Few Shot" in line:
                    loss, acc = line.split("=")[1].split("/")
                    fewshot_losses.append(float(loss))
                    fewshot_accs.append(float(acc))

        steps = list(range(len(fewshot_accs)))
        plt.cla()
        plt.xlabel("steps")
        plt.plot(steps, fewshot_accs, label="accuracy")
        plt.legend()
        plt.savefig(f"{log_dir}/{filename.strip('.log')}_Acc.png")

        plt.cla()
        plt.xlabel("steps")
        plt.plot(steps, fewshot_losses, label="loss")
        plt.legend()
        plt.savefig(f"{log_dir}/{filename.strip('.log')}_Loss.png")

LogLoss(sys.argv[1])